// SPDX-FileCopyrightText: 2022 Demerzel Solutions Limited
// SPDX-License-Identifier: LGPL-3.0-only

using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using DotNetty.Transport.Channels;
using FluentAssertions;
using Nethermind.Core.Test.Builders;
using Nethermind.Logging;
using Nethermind.Network.P2P;
using Nethermind.Network.P2P.Analyzers;
using Nethermind.Network.P2P.Messages;
using Nethermind.Network.P2P.ProtocolHandlers;
using Nethermind.Network.Rlpx;
using Nethermind.Stats.Model;
using NSubstitute;
using NUnit.Framework;

namespace Nethermind.Network.Test.P2P;

[Parallelizable(ParallelScope.Self)]
public class SessionTests
{
    private IChannel _channel;
    private IChannelHandlerContext _channelHandlerContext;
    private IPacketSender _packetSender;
    private IChannelPipeline _pipeline;

    [SetUp]
    public void SetUp()
    {
        _channel = Substitute.For<IChannel>();
        _channelHandlerContext = Substitute.For<IChannelHandlerContext>();
        _pipeline = Substitute.For<IChannelPipeline>();
        _channelHandlerContext.Channel.Returns(_channel);
        _channel.Pipeline.Returns(_pipeline);
        _pipeline.Get<ZeroPacketSplitter>().Returns(new ZeroPacketSplitter());
        _packetSender = Substitute.For<IPacketSender>();
    }

    [Test]
    public void Constructor_sets_the_values()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyB, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.That(session.RemoteNodeId, Is.EqualTo(TestItem.PublicKeyB));
        Assert.That(session.LocalPort, Is.EqualTo(30312));
        Assert.That(session.Direction, Is.EqualTo(ConnectionDirection.Out));
        Assert.That(session.SessionId, Is.Not.EqualTo(default(Guid)));
    }

    [Test]
    public void Can_set_remaining_properties()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.PingSender = Substitute.For<IPingSender>();
        Assert.That(session.PingSender, Is.Not.Null);
        session.ObsoleteRemoteNodeId = TestItem.PublicKeyC;
        Assert.That(session.ObsoleteRemoteNodeId, Is.Not.Null);
    }

    [Test]
    public void Node_can_be_retrieved_only_after_remote_known()
    {
        Session session = new(30312, _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.Throws<InvalidOperationException>(() =>
        {
            var node = session.Node;
        });

        session.Handshake(TestItem.PublicKeyA);
        session.RemoteHost = "127.0.0.1";
        session.RemotePort = 30000;
        Assert.That(session.Node, Is.Not.Null);
    }

    [Test]
    public void Throws_when_init_called_before_the_handshake()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.Throws<InvalidOperationException>(() => session.Init(5, _channelHandlerContext, _packetSender));
    }

    [Test]
    public void Raises_event_on_init()
    {
        bool wasCalled = false;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Initialized += (s, e) => wasCalled = true;

        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        Assert.That(session.P2PVersion, Is.EqualTo(5));
        Assert.That(wasCalled, Is.True);
    }

    [Test]
    public void Sets_p2p_version_on_init()
    {
        Session session = new(30312, _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyB);
        session.Init(4, _channelHandlerContext, _packetSender);
        Assert.That(session.P2PVersion, Is.EqualTo(4));
        Assert.That(session.RemoteNodeId, Is.EqualTo(TestItem.PublicKeyB));
    }

    [Test]
    public void Raises_event_on_handshake()
    {
        bool wasCalled = false;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.HandshakeComplete += (s, e) => wasCalled = true;

        session.Handshake(TestItem.PublicKeyA);
        Assert.That(wasCalled, Is.True);
    }

    [Test]
    public void Cannot_handshake_twice()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        Assert.Throws<InvalidOperationException>(() => session.Handshake(TestItem.PublicKeyA));
    }

    [Test]
    public void Cannot_handshake_after_init()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        Assert.Throws<InvalidOperationException>(() => session.Handshake(TestItem.PublicKeyA));
    }

    [Test]
    public void Cannot_enable_snappy_before_init()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.Throws<InvalidOperationException>(() => session.EnableSnappy());
        session.Handshake(TestItem.PublicKeyA);
        Assert.Throws<InvalidOperationException>(() => session.EnableSnappy());
        session.Init(5, _channelHandlerContext, _packetSender);
    }

    [Test]
    public void Can_enable_snappy()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        ZeroNettyP2PHandler handler = new(session, LimboLogs.Instance);
        _pipeline.Get<ZeroNettyP2PHandler>().Returns(handler);
        Assert.That(handler.SnappyEnabled, Is.False);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        session.EnableSnappy();
        Assert.That(handler.SnappyEnabled, Is.True);
        _pipeline.Received().Get<ZeroPacketSplitter>();
        _pipeline.Received().AddBefore(Arg.Any<string>(), Arg.Any<string>(), Arg.Any<ZeroSnappyEncoder>());
    }

    [Test]
    public void Enabling_snappy_on_disconnected_will_not_cause_trouble()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Remote, "test");
        session.EnableSnappy();
    }

    [Test]
    public async Task Adding_protocols_when_disconnecting_will_not_cause_trouble()
    {
        bool shouldStop = false;
        int i = 0;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        void addProtocol()
        {
            IProtocolHandler required = Substitute.For<IProtocolHandler>();
            required.ProtocolCode.Returns("p2p");
            required.MessageIdSpaceSize.Returns(16);
            session.AddProtocolHandler(required);
            while (!shouldStop)
            {
                IProtocolHandler protocolHandler = Substitute.For<IProtocolHandler>();
                protocolHandler.ProtocolCode.Returns("aa" + i++);
                protocolHandler.MessageIdSpaceSize.Returns(10);
                session.AddProtocolHandler(protocolHandler);
            }
        }

        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        Task task = new(addProtocol);
        task.Start();

        await Task.Delay(20);
        session.InitiateDisconnect(DisconnectReason.Other, "test");
        await Task.Delay(10);
        shouldStop = true;
    }

    [Test]
    public void Cannot_init_before_the_handshake()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.Throws<InvalidOperationException>(() => session.Init(5, _channelHandlerContext, _packetSender));
    }

    [Test]
    public void Is_closing_is_false_when_not_closing()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.That(session.IsClosing, Is.False);
        session.Handshake(TestItem.PublicKeyA);
        Assert.That(session.IsClosing, Is.False);
        session.Init(5, _channelHandlerContext, _packetSender);
        Assert.That(session.IsClosing, Is.False);
    }

    [Test]
    public void Best_state_reached_is_correct()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.That(session.BestStateReached, Is.EqualTo(SessionState.New));
        session.Handshake(TestItem.PublicKeyA);
        Assert.That(session.BestStateReached, Is.EqualTo(SessionState.HandshakeComplete));
        session.Init(5, _channelHandlerContext, _packetSender);
        Assert.That(session.BestStateReached, Is.EqualTo(SessionState.Initialized));
    }

    [Test]
    public void Cannot_dispose_unless_disconnected()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        Assert.Throws<InvalidOperationException>(() => session.Dispose());
        session.Init(5, _channelHandlerContext, _packetSender);
        Assert.Throws<InvalidOperationException>(() => session.Dispose());

        IProtocolHandler p2p = BuildHandler("p2p", 10);
        IProtocolHandler aaa = BuildHandler("aaa", 10);
        IProtocolHandler bbb = BuildHandler("bbb", 5);
        IProtocolHandler ccc = BuildHandler("ccc", 1);
        session.AddProtocolHandler(p2p);
        session.AddProtocolHandler(aaa);
        session.AddProtocolHandler(bbb);
        session.AddProtocolHandler(ccc);

        session.InitiateDisconnect(DisconnectReason.Other, "test");
        session.Dispose();

        aaa.Received().DisconnectProtocol(DisconnectReason.Other, "test");
        bbb.Received().DisconnectProtocol(DisconnectReason.Other, "test");
        ccc.Received().DisconnectProtocol(DisconnectReason.Other, "test");

        aaa.Received().Dispose();
        bbb.Received().Dispose();
        ccc.Received().Dispose();
    }

    [Test]
    public void Raises_event_on_disconnecting()
    {
        bool wasCalled = false;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Disconnecting += (s, e) => wasCalled = true;

        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        session.InitiateDisconnect(DisconnectReason.Other);
        Assert.That(wasCalled, Is.True);
    }

    [Test]
    public void Raises_event_on_disconnected()
    {
        bool wasCalled = false;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Disconnected += (s, e) => wasCalled = true;

        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Local, "test");
        Assert.That(wasCalled, Is.True);
    }

    [Test]
    public void Disconnects_after_initiating_disconnect()
    {
        bool wasCalled = false;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Disconnecting += (s, e) => wasCalled = true;

        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        session.InitiateDisconnect(DisconnectReason.Other);
        Assert.That(wasCalled, Is.True);
        Assert.That(session.IsClosing, Is.True);
    }

    [Test]
    public void Do_not_disconnects_after_initiating_disconnect_on_static_node()
    {
        bool wasCalled = false;
        Node node = new Node(TestItem.PublicKeyA, "127.0.0.1", 8545);
        node.IsStatic = true;
        Session session = new(30312, node, _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Disconnecting += (s, e) => wasCalled = true;

        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        session.InitiateDisconnect(DisconnectReason.TooManyPeers);
        Assert.That(wasCalled, Is.False);
        Assert.That(session.IsClosing, Is.False);
    }

    [Test]
    public void Error_on_channel_when_disconnecting_channels_does_not_prevent_the_event()
    {
        bool wasCalled = false;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        _channel.DisconnectAsync().Returns(Task.FromException<Exception>(new Exception()));
        session.Disconnected += (s, e) => wasCalled = true;
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Local, "test");
        Assert.That(wasCalled, Is.True);
    }

    [Test]
    public void Error_on_context_when_disconnecting_channels_does_not_prevent_the_event()
    {
        bool wasCalled = false;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        _channelHandlerContext.DisconnectAsync().Returns(Task.FromException<Exception>(new Exception()));
        session.Disconnected += (s, e) => wasCalled = true;
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Local, "test");
        Assert.That(wasCalled, Is.True);
    }

    [Test]
    public void Can_disconnect_many_times()
    {
        int wasCalledTimes = 0;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Disconnecting += (s, e) => wasCalledTimes++;

        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        session.InitiateDisconnect(DisconnectReason.Other);
        session.InitiateDisconnect(DisconnectReason.Other);
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Local, "test");
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Remote, "test");
        Assert.That(wasCalledTimes, Is.EqualTo(1));
    }

    [Test]
    public void Can_disconnect_before_init()
    {
        int wasCalledTimes = 0;
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Disconnecting += (s, e) => wasCalledTimes++;
        session.Handshake(TestItem.PublicKeyA);
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Remote, "test");
        session.Init(5, _channelHandlerContext, _packetSender);
        Assert.That(wasCalledTimes, Is.EqualTo(1));
    }

    [Test]
    public void On_incoming_sessions_can_fill_remote_id_on_handshake()
    {
        Session session = new(30312, _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyB);
        Assert.That(session.RemoteNodeId, Is.EqualTo(TestItem.PublicKeyB));
    }

    [Test]
    public void Checks_init_arguments()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        Assert.Throws<ArgumentNullException>(() => session.Init(5, null, _packetSender), "context");
        Assert.Throws<ArgumentNullException>(() => session.Init(5, _channelHandlerContext, null), "packageSender");
    }

    [Test]
    public void Can_add_and_disconnect_many_handlers()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        IProtocolHandler aaa = BuildHandler("aaa", 10);
        IProtocolHandler bbb = BuildHandler("bbb", 5);
        IProtocolHandler ccc = BuildHandler("ccc", 1);
        session.AddProtocolHandler(p2p);
        session.AddProtocolHandler(aaa);
        session.AddProtocolHandler(bbb);
        session.AddProtocolHandler(ccc);
        session.InitiateDisconnect(DisconnectReason.Other, "test");
        aaa.Received().DisconnectProtocol(DisconnectReason.Other, "test");
        bbb.Received().DisconnectProtocol(DisconnectReason.Other, "test");
        ccc.Received().DisconnectProtocol(DisconnectReason.Other, "test");
    }

    [Test]
    public void Cannot_add_handlers_before_p2p()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler aaa = BuildHandler("aaa", 10);
        Assert.Throws<InvalidOperationException>(() => session.AddProtocolHandler(aaa));
    }

    [Test]
    public void Cannot_add_handler_twice()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        IProtocolHandler p2pAgain = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);
        Assert.Throws<InvalidOperationException>(() => session.AddProtocolHandler(p2pAgain));
    }

    private IProtocolHandler BuildHandler(string code, int spaceSize)
    {
        IProtocolHandler handler = Substitute.For<IProtocolHandler>();
        handler.Name.Returns(code);
        handler.ProtocolCode.Returns(code);
        handler.MessageIdSpaceSize.Returns(spaceSize);
        return handler;
    }

    [Test]
    [NonParallelizable]
    public void Can_deliver_messages()
    {
        Metrics.P2PBytesSent = 0;

        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        IProtocolHandler aaa = BuildHandler("aaa", 10);
        IProtocolHandler bbb = BuildHandler("bbb", 5);
        IProtocolHandler ccc = BuildHandler("ccc", 1);
        session.AddProtocolHandler(p2p);
        session.AddProtocolHandler(aaa);
        session.AddProtocolHandler(bbb);
        session.AddProtocolHandler(ccc);

        var message = new TestMessage();
        _packetSender.Enqueue(message).Returns(10);
        session.DeliverMessage(message);
        _packetSender.Received().Enqueue(message);
        message.WasDisposed.Should().BeTrue();

        Metrics.P2PBytesSent.Should().Be(10);
    }

    [Test]
    public void Cannot_deliver_before_initialized()
    {
        var message = new TestMessage();
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.Throws<InvalidOperationException>(() => session.DeliverMessage(message));
        session.Handshake(TestItem.PublicKeyA);
        Assert.Throws<InvalidOperationException>(() => session.DeliverMessage(message));
        message.WasDisposed.Should().BeTrue();
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);
    }

    [Test]
    public void Cannot_receive_before_initialized()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        Assert.Throws<InvalidOperationException>(() => session.ReceiveMessage(new Packet("p2p", 1, [])));
        session.Handshake(TestItem.PublicKeyA);
        Assert.Throws<InvalidOperationException>(() => session.ReceiveMessage(new Packet("p2p", 1, [])));
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);
    }

    [Test]
    public void Stops_delivering_messages_after_disconnect()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);

        session.InitiateDisconnect(DisconnectReason.Other);

        var message = new TestMessage();
        session.DeliverMessage(message);
        _packetSender.DidNotReceive().Enqueue(Arg.Any<TestMessage>());
        message.WasDisposed.Should().BeTrue();
    }

    [Test]
    public void Stops_receiving_messages_after_disconnect()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);

        session.InitiateDisconnect(DisconnectReason.Other);

        session.ReceiveMessage(new Packet("p2p", 3, []));
        p2p.DidNotReceive().HandleMessage(Arg.Is<Packet>(static p => p.Protocol == "p2p" && p.PacketType == 3));
    }

    [Test]
    public void Delay_disconnect_until_after_initialize()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.InitiateDisconnect(DisconnectReason.TooManyPeers);

        IProtocolHandler p2p = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);

        session.Init(5, _channelHandlerContext, _packetSender);

        p2p.Received().DisconnectProtocol(Arg.Any<DisconnectReason>(), Arg.Any<string?>());
    }

    [Test]
    public void Protocol_handler_can_send_message_on_disconnect()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.InitiateDisconnect(DisconnectReason.TooManyPeers);

        IProtocolHandler p2p = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);

        var message = new TestMessage();
        p2p.When(it => it.DisconnectProtocol(Arg.Any<DisconnectReason>(), Arg.Any<string>()))
            .Do((_) =>
            {
                session.DeliverMessage(message);
            });

        session.Init(5, _channelHandlerContext, _packetSender);
        session.InitiateDisconnect(DisconnectReason.Other);

        _packetSender
            .Received()
            .Enqueue(message);

        message.WasDisposed.Should().BeTrue();
    }

    [Test, Retry(3)]
    [Parallelizable(ParallelScope.None)] // It touches global metrics
    public void Can_receive_messages()
    {
        Metrics.P2PBytesReceived = 0;

        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, NullDisconnectsAnalyzer.Instance, LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        IProtocolHandler aaa = BuildHandler("aaa", 10);
        IProtocolHandler bbb = BuildHandler("bbb", 5);
        IProtocolHandler ccc = BuildHandler("ccc", 1);
        session.AddProtocolHandler(p2p);
        session.AddProtocolHandler(aaa);
        session.AddProtocolHandler(bbb);
        session.AddProtocolHandler(ccc);

        byte[] data = new byte[10];
        session.ReceiveMessage(new Packet("---", 3, data));
        p2p.Received().HandleMessage(Arg.Is<Packet>(static p => p.Protocol == "p2p" && p.PacketType == 3));

        session.ReceiveMessage(new Packet("---", 11, data));
        aaa.Received().HandleMessage(Arg.Is<Packet>(static p => p.Protocol == "aaa" && p.PacketType == 1));

        session.ReceiveMessage(new Packet("---", 21, data));
        bbb.Received().HandleMessage(Arg.Is<Packet>(static p => p.Protocol == "bbb" && p.PacketType == 1));

        session.ReceiveMessage(new Packet("---", 25, data));
        ccc.Received().HandleMessage(Arg.Is<Packet>(static p => p.Protocol == "ccc" && p.PacketType == 0));

        session.ReceiveMessage(new Packet("---", 100, data));

        Metrics.P2PBytesReceived.Should().Be(data.Length * 5);
    }

    [Test]
    public void Updates_local_and_remote_metrics_on_disconnects()
    {
        Session session = new(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, new MetricsDisconnectsAnalyzer(), LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        IProtocolHandler p2p = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);

        long beforeLocal = Network.Metrics.LocalDisconnectsTotal.GetValueOrDefault(DisconnectReason.Other);
        long beforeRemote = Network.Metrics.RemoteDisconnectsTotal.GetValueOrDefault(DisconnectReason.Other);
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Local, string.Empty);
        long afterLocal = Network.Metrics.LocalDisconnectsTotal.GetValueOrDefault(DisconnectReason.Other);
        long afterRemote = Network.Metrics.RemoteDisconnectsTotal.GetValueOrDefault(DisconnectReason.Other);
        Assert.That(afterLocal, Is.EqualTo(beforeLocal + 1));
        Assert.That(afterRemote, Is.EqualTo(beforeRemote));

        session = new Session(30312, new Node(TestItem.PublicKeyA, "127.0.0.1", 8545), _channel, new MetricsDisconnectsAnalyzer(), LimboLogs.Instance);
        session.Handshake(TestItem.PublicKeyA);
        session.Init(5, _channelHandlerContext, _packetSender);
        p2p = BuildHandler("p2p", 10);
        session.AddProtocolHandler(p2p);

        beforeLocal = Network.Metrics.LocalDisconnectsTotal.GetValueOrDefault(DisconnectReason.Other);
        beforeRemote = Network.Metrics.RemoteDisconnectsTotal.GetValueOrDefault(DisconnectReason.Other);
        session.MarkDisconnected(DisconnectReason.Other, DisconnectType.Remote, string.Empty);
        afterLocal = Network.Metrics.LocalDisconnectsTotal.GetValueOrDefault(DisconnectReason.Other);
        afterRemote = Network.Metrics.RemoteDisconnectsTotal.GetValueOrDefault(DisconnectReason.Other);
        Assert.That(afterLocal, Is.EqualTo(beforeLocal));
        Assert.That(afterRemote, Is.EqualTo(beforeRemote + 1));
    }

    private class TestMessage : P2PMessage
    {
        public override int PacketType => P2PMessageCode.Ping;
        public override string Protocol => "p2p";

        public bool WasDisposed { get; set; }
        public override void Dispose()
        {
            base.Dispose();
            WasDisposed = true;
        }
    }
}
